package com.wz.leetcode.array;

import java.util.*;

public class TopKFrequent_347 {
    private static final Random random = new Random(System.currentTimeMillis());

    public static int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> freq = new HashMap<>();
        for (int num : nums) {
            freq.put(num, freq.getOrDefault(num, 0) + 1);
        }
        List<int[]> values = new ArrayList<>();
        for (Map.Entry<Integer, Integer> entry: freq.entrySet()) {
            int num = entry.getKey(), count = entry.getValue();
            values.add(new int[]{num, count});
        }
        int[] ret = new int[k];
        qsort(values, 0, values.size() - 1, ret, 0, k);
        return ret;
    }

    private static void qsort(List<int[]> values,
                              int start,
                              int end,
                              int[] ret,
                              int ri,
                              int k) {
        int rand = random.nextInt(end - start + 1) + start;
        Collections.swap(values, rand, end);

        // 最后一个数的频率
        int pivot = values.get(end)[1];
        int j = start - 1;
        // [start, j]都是 >= pivot的元素
        for (int i = start; i < end; i++) {
            if (values.get(i)[1] >= pivot) {
                Collections.swap(values, j + 1, i);
                j++;
            }
        }
        // 所以end 要与 j+1进行交换
        j++;
        Collections.swap(values, end, j);
        // j - start + 1个数 >= pivot
        if (k <= j - start) {
            qsort(values, start, j-1, ret, ri, k);
        } else {
            for (int i = start; i <= j; i++) {
                ret[ri++] = values.get(i)[0];
            }
            // 如果正好是k个，则不会进行递归调用，终止递归
            if (k > j - start + 1) {
                qsort(values, j + 1, end, ret, ri, k - (j - start + 1));
            }
        }
     }


    public static int[] topKFrequent1(int[] nums, int k) {
        Map<Integer, Integer> freq = new HashMap<>();
        for (int num : nums) {
            freq.put(num, freq.getOrDefault(num, 0) + 1);
        }
        PriorityQueue<int[]> minHeap = new PriorityQueue<>(k, new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return o1[1] - o2[1];
            }
        });
        int size = 0;
        for (Map.Entry<Integer, Integer> entry : freq.entrySet()) {
            int num = entry.getKey(), count = entry.getValue();

            if (size < k) {
                minHeap.offer(new int[]{num, count});
            } else {
                if (count > minHeap.peek()[1]) {
                    minHeap.poll();
                    minHeap.offer(new int[]{num, count});
                }
            }
            size++;
        }

        int[] ret = new int[k];
        for (int i = 0; i < k; i++) {
            ret[i] = minHeap.poll()[0];
        }
        return ret;
    }


    public static void main(String[] args) {
        System.out.println(Arrays.toString(topKFrequent(new int[]{1,1,1,2,2,3}, 1)));
        System.out.println(Arrays.toString(topKFrequent(new int[]{1,1,1,2,2,3}, 2)));
        System.out.println(Arrays.toString(topKFrequent(new int[]{1,1,1,2,2,3}, 3)));
        System.out.println(Arrays.toString(topKFrequent(new int[]{3,1,2,1,2,1}, 1)));
        System.out.println(Arrays.toString(topKFrequent(new int[]{3,1,2,1,2,1}, 2)));
        System.out.println(Arrays.toString(topKFrequent(new int[]{3,1,2,1,2,1}, 3)));
        System.out.println(Arrays.toString(topKFrequent(new int[]{1}, 1)));
        System.out.println("===========================================");

        System.out.println(Arrays.toString(topKFrequent1(new int[]{1,1,1,2,2,3}, 1)));
        System.out.println(Arrays.toString(topKFrequent1(new int[]{1,1,1,2,2,3}, 2)));
        System.out.println(Arrays.toString(topKFrequent1(new int[]{1,1,1,2,2,3}, 3)));
        System.out.println(Arrays.toString(topKFrequent1(new int[]{3,1,2,1,2,1}, 1)));
        System.out.println(Arrays.toString(topKFrequent1(new int[]{3,1,2,1,2,1}, 2)));
        System.out.println(Arrays.toString(topKFrequent1(new int[]{3,1,2,1,2,1}, 3)));
        System.out.println(Arrays.toString(topKFrequent1(new int[]{1}, 1)));

    }
}
